#include "MemCacheManager.h"

std::shared_ptr<void> MemCacheItem::lockMem()
{
	return m_mem;
}

std::shared_ptr<void> MemCacheItem::lockPinnedMem()
{
	if (m_pinnedMem == nullptr)
	{
		m_pinnedMem = m_manager->allocPinnedMemory(m_mem.get(), m_manager);
	}
	return m_pinnedMem;
}

std::shared_ptr<void> MemCacheItem::lockGpuMem()
{
	if (m_gpuMem == nullptr)
	{
		if (m_pinnedMem == nullptr)
		{
			m_pinnedMem = m_manager->allocPinnedMemory(m_mem.get(), m_manager);
		}
		m_gpuMem = m_manager->allocGpuMemory(m_pinnedMem.get(), m_manager);
	}
	return m_gpuMem;
}

void MemCacheItem::releasePinnedMem()
{
	m_pinnedMem = nullptr;
}

void MemCacheItem::releaseGpuMem()
{
	m_gpuMem = nullptr;
}

CoreManager::~CoreManager()
{
	//I don't know whether the locks are necessary for shared_ptr's operation
	//but I'd better lock down
	m_memCachesMutex.lock();
	m_pinnedMemCachesMutex.lock();
	m_gpuMemCachesMutex.lock();

	while (! m_memFreeCaches.empty())
	{
		auto ptr = m_memFreeCaches.top();
		m_memDeletor->free(ptr);
		m_memFreeCaches.pop();
	}
	while (! m_pinnedMemFreeCaches.empty())
	{
		auto ptr = m_pinnedMemFreeCaches.top();
		m_pinnedMemDeletor->free(ptr);
		m_pinnedMemFreeCaches.pop();
	}
	while (! m_gpuMemFreeCaches.empty())
	{
		auto ptr = m_gpuMemFreeCaches.top();
		m_gpuMemDeletor->free(ptr);
		m_gpuMemFreeCaches.pop();
	}

	m_memCachesMutex.unlock();
	m_pinnedMemCachesMutex.unlock();
	m_gpuMemCachesMutex.unlock();
}

void CoreManager::freeMem(void * ptr)
{
	m_memCachesMutex.lock();
	if (m_memCacheTotalCount >= m_memCachesSoftLimit)
	{
		m_memDeletor->free(ptr);
		m_memCacheTotalCount--;
	}
	else
	{
		m_memFreeCaches.push(ptr);
	}
	m_memCachesMutex.unlock();
}

void CoreManager::freePinnedMem(void * ptr)
{
	std::unique_lock<std::mutex> lck(m_pinnedMemCachesMutex);
	if (m_pinnedMemCacheTotalCount >= m_pinnedMemCachesSoftLimit)
	{
		m_pinnedMemDeletor->free(ptr);
		m_pinnedMemCacheTotalCount--;
	}
	else
	{
		m_pinnedMemFreeCaches.push(ptr);
		m_pinnedMemCachesCond.notify_one();
	}
	lck.unlock();
}

void CoreManager::freeGpuMem(void * ptr)
{
	std::unique_lock<std::mutex> lck(m_gpuMemCachesMutex);
	if (m_gpuMemCacheTotalCount >= m_gpuMemCachesSoftLimit)
	{
		m_gpuMemDeletor->free(ptr);
		m_gpuMemCacheTotalCount--;
	}
	else
	{
		m_gpuMemFreeCaches.push(ptr);
		m_gpuMemCachesCond.notify_one();
	}
	lck.unlock();
}

std::shared_ptr<void> CoreManager::allocPinnedMemory(const void * srcMem, std::shared_ptr<CoreManager> coreManager)
{
	//do not use lock_guard for once and all

	//try for free
	m_pinnedMemCachesMutex.lock();
	if (! m_pinnedMemFreeCaches.empty())
	{
		auto ptr = m_pinnedMemFreeCaches.top();
		m_pinnedMemFreeCaches.pop();
		m_pinnedMemCachesMutex.unlock();
		m_mem2PinnedMemMover->move(ptr, srcMem, m_memCacheBytes);
		return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_PINNED_MEM, m_pinnedMemDeletor));
	}

	//try for allocation
	if (m_pinnedMemCacheTotalCount < m_pinnedMemCachesHardLimit)
	{
		auto ptr = m_pinnedMemAllocator->alloc(m_memCacheBytes);
		if (ptr)
		{
			m_pinnedMemCacheTotalCount++;
			m_pinnedMemCachesMutex.unlock();
			m_mem2PinnedMemMover->move(ptr, srcMem, m_memCacheBytes);
			return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_PINNED_MEM, m_pinnedMemDeletor));
		}
	}
	m_pinnedMemCachesMutex.unlock();

	//try for release
	m_cacheItemsMutex.lock();
	while (m_pinnedMemFreeCaches.empty())
	{
		if (! m_cacheItems.empty())
		{
			auto item = m_cacheItems.front().lock();
			if (item)
			{
				item->releasePinnedMem();
			}
			m_cacheItems.pop();
		}
		else
		{
			break;
		}
	}
	m_cacheItemsMutex.unlock();

	std::unique_lock<std::mutex> lck(m_pinnedMemCachesMutex);
	while (m_pinnedMemFreeCaches.empty())
	{
		m_pinnedMemCachesCond.wait(lck);
	}
	auto ptr = m_pinnedMemFreeCaches.top();
	m_pinnedMemFreeCaches.pop();
	lck.unlock();
	m_mem2PinnedMemMover->move(ptr, srcMem, m_memCacheBytes);
	return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_PINNED_MEM, m_pinnedMemDeletor));
}

std::shared_ptr<void> CoreManager::allocGpuMemory(const void * srcPinnedMem, std::shared_ptr<CoreManager> coreManager)
{
	//do not use lock_guard for once and all

	//try for free
	m_gpuMemCachesMutex.lock();
	if (! m_gpuMemFreeCaches.empty())
	{
		auto ptr = m_gpuMemFreeCaches.top();
		m_gpuMemFreeCaches.pop();
		m_gpuMemCachesMutex.unlock();
		m_pinned2GpuMemMover->move(ptr, srcPinnedMem, m_memCacheBytes);
		return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_GPU_MEM, m_gpuMemDeletor));
	}

	//try for allocation
	if (m_gpuMemCacheTotalCount < m_gpuMemCachesHardLimit)
	{
		auto ptr = m_gpuMemAllocator->alloc(m_memCacheBytes);
		if (ptr)
		{
			m_gpuMemCacheTotalCount++;
			m_gpuMemCachesMutex.unlock();
			m_pinned2GpuMemMover->move(ptr, srcPinnedMem, m_memCacheBytes);
			return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_GPU_MEM, m_gpuMemDeletor));
		}
	}
	m_gpuMemCachesMutex.unlock();

	//try for release
	m_cacheItemsMutex.lock();
	while (m_gpuMemFreeCaches.empty())
	{
		if (! m_cacheItems.empty())
		{
			auto item = m_cacheItems.front().lock();
			if (item)
			{
				item->releaseGpuMem();
			}
			m_cacheItems.pop();
		}
		else
		{
			break;
		}
	}
	m_cacheItemsMutex.unlock();

	std::unique_lock<std::mutex> lck(m_gpuMemCachesMutex);
	while (m_gpuMemFreeCaches.empty())
	{
		m_gpuMemCachesCond.wait(lck);
	}
	auto ptr = m_gpuMemFreeCaches.top();
	m_gpuMemFreeCaches.pop();
	lck.unlock();
	m_pinned2GpuMemMover->move(ptr, srcPinnedMem, m_memCacheBytes);
	return std::shared_ptr<void>(ptr, Deletor(coreManager, DeletorType::DEL_TYPE_GPU_MEM, m_gpuMemDeletor));
}

std::shared_ptr<MemCacheItem> CoreManager::alloc(std::shared_ptr<CoreManager> coreManager)
{
	m_memCachesMutex.lock();
	void *memPtr;
	if(!m_memFreeCaches.empty())
	{
		memPtr = m_memFreeCaches.top();	
		m_memFreeCaches.pop();
	}
	else if(m_memCacheTotalCount < m_memCachesHardLimit)
	{
		memPtr = m_memAllocator->alloc(m_memCacheBytes);
		m_memCacheTotalCount++;
	}
	else
	{
		memPtr = nullptr;
	}
	m_memCachesMutex.unlock();

	std::shared_ptr<MemCacheItem> item = nullptr;
	if(memPtr)
	{
		m_cacheItemsMutex.lock();
		item = std::make_shared<MemCacheItem>(
			std::shared_ptr<void>(memPtr, Deletor(coreManager, DEL_TYPE_MEM, m_memDeletor)),
			coreManager);
		m_cacheItems.push(item);
		//pop some empty cacheItems
		while (m_cacheItems.front().lock() == nullptr)
		{
			m_cacheItems.pop();
		}
		m_cacheItemsMutex.unlock();
	}
	return item;
}

void CoreManager::Deletor::operator()(void * ptr)
{
	//if manager is valid, recycle mem
	//else use default deletor to free resource
	auto m = manager.lock();
	if (m)
	{
		switch (type)
		{
		case CoreManager::DEL_TYPE_MEM:
			m->freeMem(ptr);
			break;
		case CoreManager::DEL_TYPE_PINNED_MEM:
			m->freePinnedMem(ptr);
			break;
		case CoreManager::DEL_TYPE_GPU_MEM:
			m->freeGpuMem(ptr);
			break;
		default:
			break;
		}
	}
	else
	{
		fallbackDeletor->free(ptr);
	}
}
